

'''
    This program is to compute the effective number of
    codons and aminoacids in genes and genomes.
    The program requires a pickle file with all the possible aminoacids and 
    codons names: 'all_codons_and_aminos.pkl'
    
    It also requires a text file in the same format as "codonsAndAAs.txt".
    To change the name of the file, edit the script at the line:
    <===
    filename='codonsAndAAs.txt'
    ===>
    
    ------------- How to run Codes/codon_entropy.py -------------
    python Codes/codon_entropy.py [order] [prior]
    Example:
    python Codes/codon_entropy.py 0 1
    *** prior is optional. Default is one (sunset problem).
    
    Outputs:
    1. [order]_entropies.txt: [eff No. aminos] [eff. No. codons]
    one line per gene
    2. [order]_all_genome_entropy.txt: 
    same thing as before for the entire genome

'''

import pickle
import math
import sys
import numpy as np
from codonDictionaries import get_codon_to_AA_dict

def compute_all_possible_states(l, order):
    
    if order==0:
        return [[]]
    if order==1:
        return [[v] for v in l]
    
    l2=compute_all_possible_states(l, order-1)
    final_l=[]
    for i in range(len(l)):
        for j in range(len(l2)):
            final_l.append([l[i]]+l2[j])
    return final_l

def compute_all_possible_states_tuples(l,order):
    all_possible_states=compute_all_possible_states(l,order)
    return [tuple(v) for v in all_possible_states]


def old_effc(p_codon_all_original):
    p_codon_all=list(p_codon_all_original)
    sump=np.sum(p_codon_all)
    p_codon_all=[v/sump for v in p_codon_all]
    effc=0.
    for p in p_codon_all:
       effc+=p**2
    return 1./effc

def entropy_list(l):
    '''
        l lst of positive numbers
        normalizes and returns entropy
        '''
    suml=float(sum(l))
    if suml<=0:
        return 0.
    h=0.
    for v in l:
        n=(float(v)/suml)
        if n>0:
            h-=n*math.log(n,2)
    return h



def readfile(filename):
    
    '''
        returns two lists of lists
        where each list is a gene
        first has AAs on each list
        second has codons
        '''
    
    # is line_counter is even, it is codon.
    # otherwise and amino
    line_counter=0    
    all_amino_lists=[]
    all_codon_lists=[]
    
    
    for l in open(filename):
        if line_counter%2==0:
            # a list of codons (strings)
            codons=l.split()[0].split(',')
            # saves codons for analyzing the whole genome later
            if len(codons)>0:
                all_codon_lists.append(codons)
        else:
            # a list of AAs (strings)
            aminos=l.split()[0].split(',')
            # saves aminos for analyzing the whole genome later
            if len(aminos)>0:
                all_amino_lists.append(aminos)
            assert len(codons)==len(aminos), 'codon and aminos lengths dont match '+str(codons)+' '+str(aminos)
        
        line_counter+=1
        if line_counter%100==0:
            print 'line', line_counter
    
    assert len(all_amino_lists)==len(all_codon_lists)
    
    print 'sample codon->>>> ', all_codon_lists[0][0]
    
    return all_amino_lists, all_codon_lists



def update_ngram_net(codons, aminos, ngram_net, order):


    '''
        updates transitions probabilities
        for a given gene
    '''

    for pos in range(len(codons)):
        # codon unigram/diagram...
        start=[]
        if len(codons)>pos+order:
            for i in range(pos,pos+order):
                start.append(codons[i])
            if tuple(start) not in ngram_net:
                ngram_net[tuple(start)]={}
            
            assert len(start)==order
            # what comes next? 
            end=(aminos[pos+order], codons[pos+order])
            if end not in ngram_net[tuple(start)]:
                ngram_net[tuple(start)][end]=0
            ngram_net[tuple(start)][end]+=1



def compute_transitions_no_prior(codons, aminos, order):
    
    '''
        codons and aminos are lists of lists (each list is a gene)
        order 0 -> we go 0th order (() -> amino/codon)
        order 1 -> we go 1st order (codon -> amino/codon)
        order 2 -> 2nd order (codon-codon pair -> amino/codon)
        updates the network (ngrams to codons) -ngram_net-
        **** ngram_net is {codon-unigram/diagram/trigram/...: {(amino, codon):count}}
        and the ngram counts -ngram_counts-
        *** ngram_counts is {codon-unigram/...: counts}
    '''
    
    ngram_net={}
    
    for gene in range(len(codons)):
        update_ngram_net(codons[gene], aminos[gene], ngram_net, order)
    
    return ngram_net



def compute_entropy_for_given_aa(amino_codon_net, aa, all_aminos_net, prior, verbose=False):
    
    
    '''
        aa is the given aminoacid
        amino_codon_net = {aa:{codon:prob}}
    '''

    if verbose: print 'computing for ', aa
        
    p_codon_all=[]
    # compute codon probabilities
    if aa in amino_codon_net:
        for codon in all_aminos_net[aa]:
            p_codon_all.append(amino_codon_net[aa].get(codon, 0)+prior)            
            if False: print codon, amino_codon_net[aa].get(codon, 0)+prior
        if verbose: print 'effc: ', 2**entropy_list(p_codon_all), 'old', old_effc(p_codon_all), len(all_aminos_net[aa])
        # entropy for codons
        if aa=='arginine':
            print p_codon_all, 'old'
        codon_entropy_for_aa= entropy_list(p_codon_all)
    else:
        codon_entropy_for_aa= math.log(len(all_aminos_net[aa]), 2)
    
    random_entropy_for_aa=math.log(len(all_aminos_net[aa]), 2)
    #if verbose: print 'entropy for ', aa, ':', codon_entropy_for_aa
    return codon_entropy_for_aa, random_entropy_for_aa



def compute_entropy_given_state(transitions, all_aminos_net, prior, verbose=False, aa_dict_flag=False):
    
    
    '''
        transitions {(aa, codon):counts}
        given this state of the system 
        (0th order there is only one state)
        aa_dict_flag is a flag
        in case you just want the dict
        aa : k_aa (k_aa is the eff. no. codons for aa)
        the function returns
        1. the occurrences of this state
        2. entropy of the aminos
        3. entropy of the codons
        4. entropy of the codons in the random case
        5. eff. no. of codons unweighted        
    '''
    
    
    # if prior_on_codons is True, all codons get the same prior
    # and aminoacids get multiple priors if the have syn. codons
    prior_on_codons=False
    
    if verbose:
        print '==============================\n\ncompute_entropy_given_state::: prior:', prior 
        print 'transitions', transitions
    
    # for each amino, the list of possible codons
    amino_codon_net={}
    
    total_counts=0.
    for aa_codon_pair in transitions:
        aa= aa_codon_pair[0]
        codon= aa_codon_pair[1]
        if aa not in amino_codon_net:
            amino_codon_net[aa]={}
        if codon not in amino_codon_net[aa]:
            amino_codon_net[aa][codon]=0
        amino_codon_net[aa][codon]+=transitions[aa_codon_pair]
        total_counts+=transitions[aa_codon_pair]
    
    if verbose: print amino_codon_net
    if verbose: print '***************************************'
    
    #aa entropy
    H_aa=0.
    H_codon=0.
    H_codon_random=0.
    effn_codon_unweighted=0.
    effn_codon_random_unweighted=0.
    
    aa_norm=total_counts + prior * len(all_aminos_net)
    if prior_on_codons: 
        aa_norm=total_counts+prior*sum([len(all_aminos_net[aa]) for aa in all_aminos_net])
    p_aa_all=[]
    
    print 'aa_norm **** ', aa_norm

    # {aa : eff_no_codond}
    aa_dict={}
    for aa in all_aminos_net:
        # p(aa)
        if prior_on_codons: 
            p_aa=prior*len(all_aminos_net[aa])
        else:
            p_aa=prior
        if aa in amino_codon_net:
            # summing all the values of the codon dictionary
            p_aa+=sum(amino_codon_net[aa].values())
        p_aa/=aa_norm
        if verbose: print aa, p_aa
        p_aa_all.append(p_aa)
        # compute entropy for codons
        codon_entropy_aa, random_codon_entropy_aa= compute_entropy_for_given_aa(amino_codon_net, aa,\
                                                                                all_aminos_net, prior)
        #print 'aa:::', aa, codon_entropy_aa, p_aa

        
        H_codon+= p_aa * codon_entropy_aa
        H_codon_random+= p_aa * random_codon_entropy_aa
        effn_codon_unweighted+= 2**codon_entropy_aa 
        effn_codon_random_unweighted+= 2**random_codon_entropy_aa
        if aa_dict_flag:
            # skipping single codons
            if random_codon_entropy_aa>0.:
                aa_dict[aa]=(2**codon_entropy_aa, (2**codon_entropy_aa-1)/(2**random_codon_entropy_aa-1))

    if verbose: print '********', 2**H_codon, 2**H_codon_random, effn_codon_unweighted, effn_codon_random_unweighted
    
    H_aa=entropy_list(p_aa_all)
    
    if verbose: print sum(p_aa_all), 'aa_norm_check'
    assert math.fabs(sum(p_aa_all)-1)<1e-6, 'error in aa_norm ' + str(math.fabs(sum(p_aa_all)-1))    
    if verbose: print 'entropies:', H_aa, H_codon
#    assert math.fabs(effn_codon_random_unweighted-61.)<1e-6
    
    if aa_dict_flag:
        return aa_dict
    
    return total_counts+prior, H_aa, H_codon, H_codon_random, effn_codon_unweighted



def compute_entropies(codon_lists, amino_lists, all_codons, all_aminos_net, order, prior, \
                      verbose=False, roughly_random_flag=False, aa_dict_flag=False, state_all_flag=False):
    
    '''
        codon_lists is a list of genes (i.e. a list of codons)
        same thing for amino_lists
        all_codons and all_aminos_net is standard information read from pickle file
        returns the effective No. of codons weighted and unweighted
        if aa_dict_flag is True, the function returns a dict { aa: (eff no codons, compressibility) }
        if first_order_all_dict==True
        returns a dict state_all_dict = { state: compressibility }
    '''
        
    #print '***all*** ', all_codons, all_aminos_net
    if verbose: print 'order->>>', order
    
    ngram_net = compute_transitions_no_prior(codon_lists, amino_lists, order)
    #print 'ngram_net'
    #print ngram_net
    all_possible_states=compute_all_possible_states_tuples(all_codons, order)
    
    total_counts_maximum, H_aa_maximum, H_codon_maximum, \
    H_codon_random_maximum, effn_codon_unweighted_max = compute_entropy_given_state({}, \
                                                                                    all_aminos_net, \
                                                                                    prior)
    # effn_codon_unweighted_max is 61
    #assert math.fabs(effn_codon_unweighted_max-61.)<1e-6
    assert math.fabs(total_counts_maximum-prior)<1e-6
    
    if verbose: print 'maximum: ', 2**H_aa_maximum, 2**(H_codon_maximum+H_aa_maximum), total_counts_maximum
    
    # dealing with aa_dict_flag
    if aa_dict_flag:
        print 'computing aa_dict'
        assert order==0, 'aa_dict_flag only works for order zero.'
        for state in all_possible_states:
            assert state in ngram_net, 'zero order not found. This should not happen!'
            aa_dict= compute_entropy_given_state(ngram_net[state], all_aminos_net, prior, aa_dict_flag=True)
            return aa_dict
    
    
    # list of tuples (counts, aa_entropy, codon_entropy, codon_entropy_random, effn_codon_unw)
    counts_entropy=[]
    if verbose: print 'all_possible_states: ', len(all_possible_states)
    
    # { state: compressibility }
    state_all_dict={}
    
    for state in all_possible_states:
        if state in ngram_net:
            if verbose: print state, 'state'
            total_counts, H_aa, H_codon, \
            H_codon_random, effn_codon_unw= compute_entropy_given_state(ngram_net[state], \
                                                                        all_aminos_net, \
                                                                        prior)
            
            counts_entropy.append((total_counts, H_aa, H_codon, H_codon_random, effn_codon_unw))
            
            # inserting compressibility for this state
            state_all_dict[state]= 20+ 41 *(2**H_codon-1.)/(2**H_codon_random-1.)
            
        else:
            counts_entropy.append((total_counts_maximum, H_aa_maximum, H_codon_maximum, H_codon_random_maximum, effn_codon_unweighted_max))
        
    # if we want to see what happens for each state (otherwise we average among states)
    if state_all_flag:
        return state_all_dict


    # THIS IS TMP
    return counts_entropy
    print 'fix this please'
    exit()
    
    # averaging entropies and unweighted #codons of states
    if verbose: print '#states', len(counts_entropy)
    
    total_counts_states=float(sum([c[0] for c in counts_entropy]))
    if verbose: print 'total_counts_states', total_counts_states
    H_aa_states=0.
    H_codon_states=0.
    H_codon_states_random=0.
    effn_codon_unw_averaged=0.
    for c in counts_entropy:
        H_aa_states+=float(c[0])/total_counts_states * c[1]
        H_codon_states+=float(c[0])/total_counts_states * c[2]
        H_codon_states_random+=float(c[0])/total_counts_states * c[3]
        effn_codon_unw_averaged+=float(c[0])/total_counts_states * c[4]
    
    if verbose: print 2**H_aa_states, 2**(H_codon_states+H_aa_states)
    if verbose: print '########################################'

    #assert len(all_aminos_net.keys())==20
        
    #print H_codon_states, H_aa_states, 'entropies'
    compressibility=(2**H_codon_states-1.)/(2**H_codon_states_random-1.)
    compressibility_unw=(effn_codon_unw_averaged-20.)/(61.-20.)
    
    assert math.fabs(20 + compressibility_unw * (61.-20) - effn_codon_unw_averaged)<1e-6
    
    # estimate of the number of random codons (weighted)
    roughly_random = (2**H_codon_states_random)*20
    
    if verbose:
        print 'aminos', all_aminos_net.keys()
        print 'weighted compressibility', compressibility
        print 'weighted effn:',  20 + compressibility * (61.-20)
        print 'unweighted compressibility', compressibility_unw
        print 'unweighted effn:', 20 + compressibility_unw * (61.-20), 'check: ', effn_codon_unw_averaged


    if roughly_random_flag:
        return (20. + compressibility * (61.-20), effn_codon_unw_averaged, roughly_random)
    else:
        return (20. + compressibility * (61.-20), effn_codon_unw_averaged)




def run_on_file(filename, order, prior):

    # reading all codons and all aminos from pickle file
    all_codons, all_aminos_net=pickle.load(open('all_codons_and_aminos.pkl'))
    print '# codons: ', len(all_codons)
    print '# aminos: ', len(all_aminos_net)
    
    
    print all_codons, all_aminos_net, '<<<'
    # all genes
    all_amino_lists, all_codon_lists= readfile(filename)
    
    
    state_dict_ = compute_entropies(all_codon_lists, all_amino_lists, \
                            all_codons, all_aminos_net, order, prior, state_all_flag=True)
    
    
    
    state_list=[]
    for c, effc in state_dict_.iteritems():
        state_list.append([effc, c])
    print sorted(state_list)
    
    # open outfile    
    outfile=open(str(order)+'_entropies.txt', 'w')
    for gene_index in range(len(all_codon_lists)):
        # computing model and entropies for each gene
        if gene_index%100==0:
            print gene_index, 'genes done...'
        assert len(all_codon_lists[gene_index])==len(all_amino_lists[gene_index])
        effn_codons_w, effn_codons_uw, roughly_random= compute_entropies([all_codon_lists[gene_index]], \
                                                          [all_amino_lists[gene_index]], \
                                                          all_codons, all_aminos_net, order, prior, roughly_random_flag=True)
        outfile.write(str(len(all_codon_lists[gene_index]))+' '+str(effn_codons_w)+\
                      ' '+str(effn_codons_uw)+' '+str(roughly_random)+'\n')
    outfile.close()
    print len(all_codon_lists), 'genes done. Now all genome...'
    
    # for the entire genome
    outfile=open(str(order)+'_all_genome_entropy.txt', 'w')
    effn_codons_w, effn_codons_uw, roughly_random = compute_entropies(all_codon_lists, all_amino_lists, \
                                                      all_codons, all_aminos_net, order, prior,roughly_random_flag=True)
    outfile.write(str(effn_codons_w)+' '+str(effn_codons_uw)+' '+str(roughly_random)+'\n')
    outfile.close()
        
    print 'done!'
    print 'computing aa_dict for the entire genome'
    print compute_entropies(all_codon_lists, all_amino_lists, \
                            all_codons, all_aminos_net, 0, prior, aa_dict_flag=True)
                            
                            
            
def run_on_file_simple_codon_entropy(filename, prior_aa=1.):
    
    
    '''
        computing the entropy of entire genome
        assuming prior_aa on the aminos
    '''
    
    # reading all codons and all aminos from pickle file
    all_codons, all_aminos_net=pickle.load(open('all_codons_and_aminos.pkl'))
    print '# codons: ', len(all_codons)
    print '# aminos: ', len(all_aminos_net)
    
    
    #print all_codons, all_aminos_net, '<<<'
    # all genes
    all_amino_lists, all_codon_lists= readfile(filename)
    
    codon_to_AA= get_codon_to_AA_dict()
    
    # { codon : occurrences }
    codon_hist={}
    for gene in all_codon_lists:
        for c in gene:
            codon_hist[c]=codon_hist.get(c,0)+1
    
    # { aa : occurrences }
    aa_hist={}
    for gene in all_codon_lists:
        for c in gene:
            aa_hist[codon_to_AA[c]]=aa_hist.get(codon_to_AA[c],prior_aa)+1

    assert len(aa_hist) == len(all_aminos_net), 'fix this...'
    
    
    # now compute the entropy of codon given aa
    # { aa : {codon: occs} }
    aa_codon_given_aa = {}
    for gene in all_codon_lists:
        for c in gene:
            aa= codon_to_AA[c]
            if aa not in aa_codon_given_aa:
                aa_codon_given_aa[aa]={}
            aa_codon_given_aa[aa][c]=aa_codon_given_aa[aa].get(c, prior_aa)+1
            
    
    print '>>>>>>>>>> aa_codon_given_aa'
    print aa_codon_given_aa
    
    entropy_codon_given_aa=0.
    all_aas=sum(aa_hist.values())
    
    print 'all_aas **** ', all_aas
    for aa in aa_codon_given_aa:
        p_aa= float(aa_hist[aa])/all_aas
        entropy_codon_given_aa+=p_aa * entropy_list(aa_codon_given_aa[aa].values())
        print 'aa ..', aa, entropy_list(aa_codon_given_aa[aa].values()), p_aa
        if aa=='arginine':
            print aa_codon_given_aa[aa].values(), 'new'
    
    print '>>> entropy codons given aa', entropy_codon_given_aa
    # now adding prior to codons
    for c in all_codons:
        assert c in codon_to_AA
        prior=float(prior_aa)/ len(all_aminos_net[codon_to_AA[c]])
        #print c, prior
        codon_hist[c]=codon_hist.get(c,0)+prior
        #codon_hist[c]=codon_hist.get(c,0)+1
    
    entropy_codons= entropy_list(codon_hist.values())
    print 'entropy_codons', entropy_codons, 2**entropy_codons

    entropy_aas= entropy_list(aa_hist.values())
    print 'entropy_aas', entropy_aas, 2**entropy_aas

    
    print 'check entropy given aa', entropy_codons-entropy_aas
    
    
    
    # now on the entire previous thing
    counts_entropy=compute_entropies(all_codon_lists, all_amino_lists, \
                      all_codons, all_aminos_net, 0, prior_aa)
    
    print counts_entropy[0][1]+counts_entropy[0][2], '<<<<', counts_entropy


if __name__=='__main__':
    
    if len(sys.argv)<2:
        print 'python ', sys.argv[0], '[max_order] [prior (1 is default)]'
        exit()
    
    maxorder=int(sys.argv[1])
    assert maxorder>=0, 'maxorder should be >=0'
    prior_main=1.
    if len(sys.argv)>2:
        prior_main=float(sys.argv[2])
    
    print 'prior: ', prior_main
    #filename='smallC.txt'
    filename='escherichia.txt'
    
    #run_on_file_simple_codon_entropy(filename, prior_main)
    #exit()
    for order_main in range(0,maxorder+1):
        run_on_file(filename, order_main, prior_main)



